package com.jd.genie.agent.tool.common;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.jd.genie.agent.agent.AgentContext;
import com.jd.genie.agent.dto.CodeInterpreterRequest;
import com.jd.genie.agent.dto.CodeInterpreterResponse;
import com.jd.genie.agent.dto.File;
import com.jd.genie.agent.printer.Printer;
import com.jd.genie.agent.tool.BaseTool;
import com.jd.genie.agent.util.SpringContextHolder;
import com.jd.genie.agent.util.StringUtil;
import com.jd.genie.config.GenieConfig;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import org.springframework.context.ApplicationContext;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.*;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

@Slf4j
public class ReportTool implements BaseTool {

    private final AgentContext agentContext;

    private final Printer printer;

    public ReportTool(AgentContext agentContext, Printer printer)
    {
        this.agentContext = agentContext;
        this.printer = printer;
    }

    @Override
    public String getName()
    {
        return "report_tool";
    }

    @Override
    public String getDescription()
    {
        String desc = "这是一个报告工具，可以通过编写HTML、MarkDown报告";
        return """
                这是一个专业的Markdown、PPT和HTML的生成工具，可以用于输出 html 格式的网页报告、PPT 或者 markdown 格式的报告，
                生成报告或者需要输出 Markdown 或者 HTML 或者 PPT 格式时，一定使用此工具生成报告。
                （如果没有明确的格式要求默认生成 Markdown 格式的），不要重复生成相同、类似的报告。这不是查数工具，
                严禁用此工具查询数据。不同入参之间都应该是跟任务强相关的，同时文件名称、
                文件描述和任务描述之间都应该是围绕完成任务目标生成的。""";
    }

    /**
     * reportTool 配置
     */
    private final Map<String, Object> reportToolParams;

    {
        String jsonStr = """
                {
                  "type" : "object",
                  "properties" : {
                    "fileDescription" : {
                      "description" : "生成报告的文件描述，一定是跟用户的任务强相关的",
                      "type" : "string"
                    },
                    "fileName" : {
                      "description" : "生成的文件名称，文件的前缀中文名称一定是跟用户的任务和文件内容强相关，如果是markdown，则文件名称后缀是 .md，如果是ppt、html文件，则是文件名称后缀是 .html，一定要包含文件后缀。文件名称不能使用特殊符号，不能使用、，？等符号，如果需要，可以使用下划线_。",
                      "type" : "string"
                    },
                    "task" : {
                      "description" : "生成文件任务的具体要求及详细描述，以及需要在报告中体现的内容，例如，上下文中需要输出的数据细节。",
                      "type" : "string"
                    },
                    "fileType" : {
                      "description" : "仅支持 markdown ppt 和 html 三种类型，如果指定了输出 html 或 网页版 格式，则是html，如果指定了输出 ppt、pptx 则是 ppt，否则使用 markdown 。",
                      "type" : "string"
                    }
                  },
                  "required" : [ "fileType", "task", "fileName", "fileDescription" ]
                }
                """;
        reportToolParams = JSON.parseObject(jsonStr, Map.class);
    }

    @Override
    public Map<String, Object> toParams()
    {
        if (!reportToolParams.isEmpty()) {
            return reportToolParams;
        }

        Map<String, Object> taskParam = new HashMap<>();
        taskParam.put("type", "string");
        taskParam.put("description", "需要完成的任务以及完成任务需要的数据，需要尽可能详细");
        Map<String, Object> parameters = new HashMap<>();
        parameters.put("type", "object");
        Map<String, Object> properties = new HashMap<>();
        properties.put("task", taskParam);
        parameters.put("properties", properties);
        parameters.put("required", Collections.singletonList("task"));

        return parameters;
    }

    @Override
    public Object execute(Object input)
    {
        try {
            Map<String, Object> params = (Map<String, Object>) input;
            String task = (String) params.get("task");
            String fileDescription = (String) params.get("fileDescription");
            String fileName = (String) params.get("fileName");
            String fileType = (String) params.get("fileType");

            if (fileName.isEmpty()) {
                String errMessage = "文件名参数为空，无法生成报告。";
                log.error("{} {}", agentContext.getRequestId(), errMessage);
                return null;
            }

            List<String> fileNames = agentContext.getProductFiles().stream().map(File::getFileName).collect(Collectors.toList());
            Map<String, Object> streamMode = new HashMap<>();
            streamMode.put("mode", "token");
            streamMode.put("token", 10);
            CodeInterpreterRequest request = CodeInterpreterRequest.builder()
                    .requestId(agentContext.getSessionId()) // 适配多轮对话
                    .query(agentContext.getQuery())
                    .task(task)
                    .fileNames(fileNames)
                    .fileName(fileName)
                    .fileDescription(fileDescription)
                    .stream(true)
                    .contentStream(agentContext.getIsStream())
                    .streamMode(streamMode)
                    .fileType(fileType)
                    .build();
            return callCodeAgentStream(request);
        } catch (Exception e) {
            log.error("{} report_tool error", agentContext.getRequestId(), e);

        }
        return null;
    }

    /**
     * 调用 CodeAgent
     */
    @SneakyThrows
    public String callCodeAgentStream(CodeInterpreterRequest codeRequest)
    {
        OkHttpClient client = new OkHttpClient.Builder()
                .connectTimeout(60, TimeUnit.SECONDS) // 设置连接超时时间为 1 分钟
                .readTimeout(600, TimeUnit.SECONDS)    // 设置读取超时时间为 10 分钟
                .writeTimeout(600, TimeUnit.SECONDS)   // 设置写入超时时间为 10 分钟
                .callTimeout(600, TimeUnit.SECONDS)    // 设置调用超时时间为 10 分钟
                .build();

        ApplicationContext applicationContext = SpringContextHolder.getApplicationContext();
        GenieConfig genieConfig = applicationContext.getBean(GenieConfig.class);
        String url = genieConfig.getCodeInterpreterUrl() + "/v1/tool/report";
        RequestBody body = RequestBody.create(
                MediaType.parse("application/json"),
                JSONObject.toJSONString(codeRequest)
        );

        log.info("{} report_tool request {}", agentContext.getRequestId(), JSONObject.toJSONString(codeRequest));
        Request.Builder requestBuilder = new Request.Builder()
                .url(url)
                .post(body);
        Request request = requestBuilder.build();

        final int firstInterval = 1;
        final int sendInterval = 4;
        try (Response response = client.newCall(request).execute();
             ResponseBody responseBody = response.body()) {
            log.info("{} report_tool response {} {} {}", agentContext.getRequestId(), response, response.code(), response.body());

            if (!response.isSuccessful() || responseBody == null) {
                log.error("{} report_tool request error.", agentContext.getRequestId());
                throw new IOException("Unexpected response code: " + response);
            }
            CodeInterpreterResponse codeResponse = CodeInterpreterResponse.builder()
                    .codeOutput("report_tool 执行失败") // 默认输出
                    .build();

            int index = 1;
            StringBuilder stringBuilderIncr = new StringBuilder();
            String line;
            String messageId = StringUtil.getUUID();
            // 获取数字人名称
            String digitalEmployee = agentContext.getToolCollection().getDigitalEmployee(getName());
            BufferedReader reader = new BufferedReader(new InputStreamReader(responseBody.byteStream()));
            while ((line = reader.readLine()) != null) {
                if (!line.startsWith("data: ")) {
                    continue;
                }
                String data = line.substring(6);
                if (data.equals("[DONE]")) {
                    break;
                }
                if (index == 1 || index % 100 == 0) {
                    log.info("{} report_tool recv data: {}", agentContext.getRequestId(), data);
                }
                if (data.startsWith("heartbeat")) {
                    continue;
                }
                codeResponse = JSONObject.parseObject(data, CodeInterpreterResponse.class);
                if (codeResponse.getIsFinal()) {
                    // report_tool 只会输出一个文件，使用模型输出的文件名和描述
                    if (Objects.nonNull(codeResponse.getFileInfo())) {
                        for (CodeInterpreterResponse.FileInfo fileInfo : codeResponse.getFileInfo()) {
                            File file = File.builder()
                                    .fileName(codeRequest.getFileName())
                                    .fileSize(fileInfo.getFileSize())
                                    .ossUrl(fileInfo.getOssUrl())
                                    .domainUrl(fileInfo.getDomainUrl())
                                    .description(codeRequest.getFileDescription())
                                    .isInternalFile(false)
                                    .build();
                            agentContext.getProductFiles().add(file);
                            agentContext.getTaskProductFiles().add(file);
                        }
                    }
                    printer.send(messageId, codeRequest.getFileType(), codeResponse, digitalEmployee, true);
                } else {
                    stringBuilderIncr.append(codeResponse.getData());
                    if (index == firstInterval || index % sendInterval == 0) {
                        codeResponse.setData(stringBuilderIncr.toString());
                        printer.send(messageId, codeRequest.getFileType(), codeResponse, digitalEmployee, false);
                        stringBuilderIncr.setLength(0);
                    }
                }
                index++;
            }
            // 统一使用data字段，兼容历史codeOutput逻辑
            return Objects.nonNull(codeResponse.getData()) && !codeResponse.getData().isEmpty() ? codeResponse.getData() : codeResponse.getCodeOutput();
        }
    }

}